from langchain_community.llms import Ollama
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.messages import SystemMessage, HumanMessage

output_parser = StrOutputParser()

# 直接在创建 Ollama 对象时传递参数,设置远程地址
llm = Ollama(model="deepseek-r1:latest", temperature=0,top_p = 0.,base_url="http://192.168.13.201:11434")



# 创建一个聊天prompt模板
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一位世界级的技术文档写手，你的名字是{name}"),
    ("user", "{input}")
])

# prompt = ChatPromptTemplate.from_messages([
#     SystemMessage(content="你是一位世界级的技术文档写手，你的名字是{name}"),
#     HumanMessage(content="{input}")
# ])


#创建一个简单prompt
# prompt = ChatPromptTemplate.from_template("你是一位世界级的技术文档写手，你的名字是{name}，请回答我的问题：{input}")

#通过prompt模板和llm模型，构建一个chain
chain = prompt | llm | output_parser

#实现用户输入一次，然后回答一次
while True:
    user_input = input("请输入：")

   
    
    #使用 stream 方法实现流式输出
    for chunk in llm.stream(prompt.format(name="贾维斯",input=user_input)):
        print(chunk, end="", flush=True)
    print("\n")  # 添加换行

